# -*- coding: utf-8 -*-
"""
Created on Tue May  7 13:25:21 2019

@author: Soly Liang
"""
import numpy as np

def AnalysisData(X):
    # 分析数据
    # DataType 返回每个属性的数据类型，若为字符串类型，说明是离散型，并统计各个可能取值
    # X1 返回数字化的数据，将离散型数值转变为相应取值的编号
    m=len(X)     #样本数
    n=len(X[0])  #属性数
    DataType=[]
    for p in range(n):
        DataType.append([type(X[0][p])])
        if DataType[p][0]==str:
            for q in range(m):
                if X[q][p] not in DataType[p]:
                    DataType[p].append(X[q][p])
    X1=[]   #数字化的数据
    for p in range(m):
        temp=[]
        for q in range(n):
            if DataType[q][0]==str:
                temp.append(DataType[q].index(X[p][q]))
            else:
                temp.append(X[p][q])
        X1.append(temp)
    return DataType,X1

def most(Y):
    # 统计样本集中的多数类
    Yvalue=[]  #Y取值
    Ynum=[]    #取值数目
    for y in Y:
        if y not in Yvalue:
            Yvalue.append(y)
            Ynum.append(1)
        else:
            Ynum[Yvalue.index(y)]+=1
    return Yvalue[Ynum.index(max(Ynum))]

def Ent(Y):
    # 计算信息熵
    if Y==[]:
        return 0
    Yvalue=[]  #Y取值
    Ynum=[]    #取值数目
    for y in Y:
        if y not in Yvalue:
            Yvalue.append(y)
            Ynum.append(1)
        else:
            Ynum[Yvalue.index(y)]+=1
    pk=np.array(Ynum)/sum(Ynum)
    return sum(-pk*np.log2(pk))

def choose(X,Y,FeatureName,considerfeature,DataType):
    Gain=[] #用于存储信息增益值
    for f in considerfeature:
        if DataType[f][0]==str:  #若该特征取值为离散型
            Gainf=Ent(Y)       
            for v in range(1,len(DataType[f])):
                Yv=Y[X[:,f]==v]
                Gainf-=Ent(Yv)*len(Yv)/len(Y)
            Gain.append([Gainf])
        else:   #若该特征取值为连续型
            Gainf=[]
            values=np.array(list(set(X[:,f])))  #利用集合去除重复取值
            t=(np.sort(values)[:-1]+np.sort(values)[1:])/2  #候选划分点
            for tt in t:
                Y1=Y[X[:,f]<=tt]; Y2=Y[X[:,f]>tt]
                Gaint=Ent(Y)-Ent(Y1)*len(Y1)/len(Y)-Ent(Y2)*len(Y2)/len(Y)
                Gainf.append(Gaint)
            Gain.append([max(Gainf),t[Gainf.index(max(Gainf))]])
    GainValue=[g[0] for g in Gain]
    maxGain=max(GainValue)
    index=GainValue.index(maxGain)
    '''
    #当多个属性值对应的信息增益相等时，提示选择其中一个
    if GainValue.count(maxGain)!=1:
        index=np.array(range(len(GainValue)))[GainValue==maxGain]
        print('现有%d个属性所得信息增益相等,请从中选择一个：'%len(index))
        k=1
        for i in index:
            print('%d:%s'%(k,FeatureName[considerfeature[i]]),end=', ')
            k+=1
        n=input('')
        index=index[int(n)-1]
    else:
        index=GainValue.index(maxGain)
    '''
    if len(Gain[index])>1:  #返回最佳属性编号，若为连续取值，还要返回划分点
        return [considerfeature[index],Gain[index][1]]
    else:
        return [considerfeature[index]]        

def CreatTree(X,Y,FeatureName,LabelName,considerfeature='all',DataType=None):
    if considerfeature=='all':   #初始时的'all'表示考虑所有特征
        considerfeature=list(range(len(FeatureName)))
    if DataType==None:           #初始时需要分析数据，将X数字化，并得到每个X的数据类型
        DataType,X=AnalysisData(X)
        X=np.array(X)
        Y=np.array(Y)
    if (Y==Y[0]).all():      #若所有样本属于同一类
        return LabelName[Y[0]]
    if len(considerfeature)==0:  #若待考察特征为空
        return LabelName[most(Y)]
    if (X[:,considerfeature]==X[0,considerfeature]).all():  #若所有样本取值相同
        return LabelName[most(Y)] 
    BestFeature=choose(X,Y,FeatureName,considerfeature,DataType)      #选择最佳划分属性
    bf=BestFeature[0]
    #tree={FeatureName[bf]:{}}
    if len(BestFeature)==1:  #若为离散型属性
        tree={FeatureName[bf]+'=?':{}}
        considerfeature.remove(bf)
        for v in range(1,len(DataType[bf])):
            Xv=X[X[:,bf]==v,:]
            Yv=Y[X[:,bf]==v]
            if len(Yv)==0:  #若该分支的训练集为空
                tree[FeatureName[bf]+'=?'][DataType[bf][v]]=LabelName[most(Y)]
            else:
                # ↓↓这一句特别注意，列表传入函数后，若函数中对列表作了操作
                # 函数外的这个列表也将随之改变，为了避免这种情况，可以用list.copy()功能
                tree[FeatureName[bf]+'=?'][DataType[bf][v]]=CreatTree(Xv,Yv,FeatureName,LabelName,considerfeature.copy(),DataType)
    else: #若为连续型属性
        tree={FeatureName[bf]:{}}
        t=BestFeature[1]  #划分点
        X1=X[X[:,bf]<=t,:]; Y1=Y[X[:,bf]<=t]
        X2=X[X[:,bf]>t,:];  Y2=Y[X[:,bf]>t]
        tree[FeatureName[bf]]['<=%.3f'%t]=CreatTree(X1,Y1,FeatureName,LabelName,considerfeature.copy(),DataType)
        tree[FeatureName[bf]]['>%.3f'%t]=CreatTree(X2,Y2,FeatureName,LabelName,considerfeature.copy(),DataType)
    return tree